#!/usr/bin/python
'''
Evaluate the mapping results.
'''
import sys
import os
import gzip

# Parameters.
read_file = sys.argv[1]
map_file = sys.argv[2]

def read_file(filepath):
	# define data structure.
	data = {}
	
	# Check if zipped.
	zipped = False
	if filepath.count(".gz"):
		zipped = True
	
	# Open file.
	seq = ""
	seqheader = ""
	first = True
	
	if zipped == True:
		fin = gzip.open(filepath,"r")
	else:
		fin = open(filepath,"r")
		

	for line in fin:
		
		# Skip comments.	
		if line[0] == "#": continue

		# first time.
		if first == True and line[0] == ">":
			seq = ""
			seqheader = line.strip().replace(">","")
			first = False
			continue		

		# finished seq.
		if line[0] == ">":
			# save data.
			data[seqheader] = seq.strip()

			# Reset variables.
			seq = ""
			seqheader = line.strip().replace(">","")
			continue

		# get seq.
		seq += line.strip()
			
	# Handle last close and return.
	data[seqheader] = seq.strip()
	fin.close()
	return data


# Read use list.

print "Reading names."
fin = open(sys.argv[2], "rb")
names = fin.readlines()
fin.close()
print "done"


# Read in fastq.
print "reading seqs"
seqs = read_file(sys.argv[1])
print "done", len(seqs)

# Only print ones in names.
fout = open(sys.argv[3],"wb")
cnt = 0
for name in names:
	name = name.strip().replace(">","")
	if name in seqs:
		cnt += 1
		if cnt % 1000000 == 0:
			print cnt
			fout.flush()
		
		fout.write(">%s\n%s\n" % (name, seqs[name]))
fout.close()
